#!/usr/bin/env python
import dpdata
import matplotlib.pyplot as plt
import numpy as np

# 读取训练数据点
training_systems = dpdata.LabeledSystem("/home/zxg/be-cu_dp_course/dpgen/run/deepmd_file/Be16_hcp0.96_1.04/deepmd", fmt = "deepmd/npy")
predict = training_systems.predict("/home/zxg/be-cu_dp_course/dpgen/run/iter.000000/00.train/000/frozen_model.pb")

# 提取训练和预测的力分量
train_forces = training_systems["forces"]
predict_forces = predict["forces"]

# 提取 x、y、z 轴力的分量
train_forces_x = [force[0] for sublist in train_forces for force in sublist]
train_forces_y = [force[1] for sublist in train_forces for force in sublist]
train_forces_z = [force[2] for sublist in train_forces for force in sublist]

predict_forces_x = [force[0] for sublist in predict_forces for force in sublist]
predict_forces_y = [force[1] for sublist in predict_forces for force in sublist]
predict_forces_z = [force[2] for sublist in predict_forces for force in sublist]

# 绘制散点图
plt.figure(figsize=(12, 9))  # 设置figure的大小

plt.subplot(2, 2, 1)
plt.scatter(training_systems["energies"], predict["energies"], label='energy', color='blue', marker='o')
plt.xlabel('Energy of DFT')
plt.ylabel('Energy of DeepMD')

# 绘制 x 轴力的散点图
plt.subplot(2, 2, 2)
plt.scatter(train_forces_x, predict_forces_x, label='X Forces', color='blue', marker='o')
plt.xlabel('Forces x of DFT')
plt.ylabel('Forces x of DeepMD')

# 绘制 y 轴力的散点图
plt.subplot(2, 2, 3)
plt.scatter(train_forces_y, predict_forces_y, label='Y Forces', color='green', marker='x')
plt.xlabel('Forces y of DFT')
plt.ylabel('Forces y of DeepMD')

# 绘制 z 轴力的散点图
plt.subplot(2, 2, 4)
plt.scatter(train_forces_z, predict_forces_z, label='Z Forces', color='red', marker='^')
plt.xlabel('Forces z of DFT')
plt.ylabel('Forces z of DeepMD')

x_range = np.linspace(plt.xlim()[0], plt.xlim()[1])

plt.plot(x_range, x_range, "r--", linewidth = 0.25)
# 添加图例
plt.legend()

# 设置图片的分辨率和保存路径
plt.savefig('inter00_f_e_diff_test.png', dpi=300)